#ifndef __MEMORY_MANAGER_HPP
#define __MEMORY_MANAGER_HPP

#include <memory>
#include <benchmark/benchmark.h>

class CustomMemoryManager : public benchmark::MemoryManager
{
public:
    int64_t num_allocs;
    int64_t max_bytes_used;

    void Start() BENCHMARK_OVERRIDE
    {
        num_allocs = 0;
        max_bytes_used = 0;
    }

    void Stop(Result *result) BENCHMARK_OVERRIDE
    {
        result->num_allocs = num_allocs;
        result->max_bytes_used = max_bytes_used;
    }
};

std::unique_ptr<CustomMemoryManager> mm(new CustomMemoryManager());

#ifdef MEMORY_PROFILER
void *custom_malloc(size_t size)
{
    void *p = malloc(size);

    auto *mm_ptr = mm.get();
    if (mm_ptr) {
        mm_ptr->num_allocs += 1;
        mm_ptr->max_bytes_used += size;
    }
    
    return p;
}

#define malloc(size) custom_malloc(size)

void *operator new(size_t size)
{
    void *p = malloc(size);
    return p;
}

void operator delete(void *p) noexcept
{
    free(p);
}

void *operator new[](size_t size)
{
    void *p = malloc(size);
    return p;
}

void operator delete[](void *p) noexcept
{
    free(p);
}
#endif

#endif // __MEMORY_MANAGER_HPP